package com.kfm.kfmchat.listener;

import cn.hutool.core.util.ObjectUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.json.JsonMapper;
import com.unfbx.chatgpt.entity.chat.ChatChoice;
import com.unfbx.chatgpt.entity.chat.ChatCompletionResponse;
import com.unfbx.chatgpt.entity.chat.Message;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.util.ObjectUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.awt.*;
import java.io.IOException;

// 只要继承了 EventSourceListener 就可以监听 EventSource 的变化
@Slf4j
@Setter
public class ChatStreamEventSourceListener extends EventSourceListener {

    private SseEmitter sseEmitter;
    private StringBuilder allData;

    @Override
    public void onClosed(@NotNull EventSource eventSource) {
        log.info("关闭");
        sseEmitter.complete();
    }

    @Override
    public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {
        log.info("获取到响应的内容{}", data);
        // 流数据响应结束
        if (data.isEmpty() || ObjectUtil.contains(data, "[DONE]")) {
            try {
                sseEmitter.send(SseEmitter.event().data("DONE").id("1").name("message"));
                sseEmitter.complete();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            return;
        }
        // data [DONE]
        ChatCompletionResponse response = null;
        try {
            response = new JsonMapper().readValue(data, ChatCompletionResponse.class);
        } catch (JsonProcessingException e) {
            throw new RuntimeException(e);
        }

        ChatChoice chatChoice = response.getChoices().get(0);
        // TODO 如果响应结束 将所有响应的内容拼接到一起
        if (!"stop".equals(chatChoice.getFinishReason()) && !"length".equals(chatChoice.getFinishReason())) {
            allData.append(chatChoice.getDelta().getContent());
        }
        try {
            String content = chatChoice.getDelta().getContent();
            if (content == null) {
               content = "DONE";
            }
            sseEmitter.send(SseEmitter.event().data(content).id("1").name("message"));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

    }

    @Override
    public void onFailure(@NotNull EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
        log.error("失败{}", t);
    }

    @Override
    public void onOpen(@NotNull EventSource eventSource, @NotNull Response response) {
        log.info("打开");
    }
}
